# Additional tools for PyCuda function definitions

"""

Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

from pycuda.compiler import SourceModule
import numpy as N

# Grid convenience functions
full_grid = lambda rows, cols, bsize: (cols/bsize, rows/bsize)
vert_grid = lambda rows, cols, bsize: (1, rows/bsize)
sing_grid = lambda rows, cols, bsize: (1,1)


# Thread indexing for a fully-gridded matrix. Requires cols
FULL_GRID_IDX = """
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int base_idx = __mul24(blockIdx.x, BLOCK_SIZE) + __mul24(blockIdx.y, BLOCK_SIZE) * cols; // pos of t(0,0) in a
    int glob_idx = base_idx + tx + cols * ty;                                                // pos of t in a

    """

# Thread indexing for a vertically-gridded (i.e., grid=(1, a.shape[0] / BLOCK_SIZE)). Requires cols
VERT_GRID_IDX = """
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int base_idx = __mul24(blockIdx.y, BLOCK_SIZE) * cols; // pos of t(0,0) in a
    int glob_idx = base_idx + tx + cols * ty;              // pos of t in a

    """

# Thread indexing for singly-gridded (i.e., grid=(1,1))
SINGLE_GRID_IDX = """
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int glob_idx = tx + cols * ty;// pos of t in a

    """

def nearest(x, mult):
    """Returns the nearest multiple of mult to x"""

    leftover = x % mult

    if leftover:
        return x + mult - leftover

    return x

def aligntobsize(M, bsize = 16):
    """Rather crude way to align M shape to a multiple of bsize by padding with zeros"""

    assert len(M.shape) == 2
    assert M.dtype == N.float32

    padding = bsize - N.mod(M.shape, bsize)

    if padding[0] != bsize:
        M = N.vstack((M, N.zeros((padding[0], M.shape[1]), M.dtype)))

    if padding[1] != bsize:
        M = N.hstack((M, N.zeros((M.shape[0], padding[1]), M.dtype)))

    return M.copy()

def unroll(loop, var, end, start = 0, step = 1):
    """
    Simple loop unroller

    Designed to replace:
    for (k = start; k < end; k += step)

    loop  = The loop string, sans for and bracing
    var   = the variable in the string to be replaced by numbers
    end   = the last digit plus 1. NOTE: AS ABOVE, THIS FUNCTION ASSUMES VAR < END!
    start = the first digit
    step  = the stepping
    """
    
    unrolled = ''

    for i in xrange(start,end,step):
        unrolled += loop.replace(var, str(i)) + '\n'

    return unrolled


class Krnl(object):
    """
    Class for retrieving prepared PyCuda function calls from simple string kernels

    Krnl(source, bsize, grid_idx, prep, fmt = {}, fname = None, **kwds)
    
    Where:
        source = The source code, with %(str) format blocks as detailed below
        bsize  = The blocksize, which is called as (bsize, bsize, 1)

                 The source code may have %(bsize)s in its definitions, e.g.:
                    #define BLOCK_SIZE %(bsize)s
                 
                 Which will be automatically replaced.

                 Can also be called as a tuple, which will be called as (bsize[0], bsize[1], 1)
                 No replacement occurs in this case.
        
        grid_idx = The indices of each thread, based on how the function is gridded.  (FULL_GRID_IDX, etc)

                   The source code should have %(GRID_IDX)s near the top of the function

        frmt    = Additional formatting beyond bsize and grid_idx to be included. Should be a dict.
        

        Additional kwds will be passed on to SourceModule
    """

    def __init__(self, source, bsize, grid_idx = '', frmt = {}, **kwds):

        if isinstance(bsize, tuple):
            source = self._krnl_str(source, '', grid_idx, frmt)
        else:
            source = self._krnl_str(source, bsize, grid_idx, frmt)
    
        self.mod   = SourceModule(source, **kwds)
        self.bsize = bsize

    def func(self, fname, prep, texrefs=[]):
        """
        prep    = The PyCuda function preparation string
        fname   = The function name in the source to retrieve
        texrefs = Texture references 
        """

        trefs = []
        bsize = self.bsize

        func = self.mod.get_function(fname)

        if isinstance(bsize, tuple):
            func.prepare(prep, texrefs=texrefs)
            block=(bsize[0], bsize[1], 1)
        else:
            func.prepare(prep, texrefs=texrefs)
            block=(bsize,bsize,1)

        #TODO: eval self.fname = func?
        
        return func.prepared_call, block

    def _krnl_str(self, source, bsize, grid_idx, frmt = {}):
    
        frmdict = {'bsize'   : bsize,
                  'GRID_IDX': grid_idx}
    
        frmdict.update(frmt)
    
        try:
            return source % frmdict
        except:
            raise ValueError, "Incorrect format string. Did you remember to escape your %% signs?"

    
elemwise_inplace_kernel = """
    #define BLOCK_SIZE %(bsize)s
    #define elem tmp[ty][tx]

    __global__ void elemwise_inplace(float *a, int cols, float val) {
      
      // cols is the width of a
      // This function affects a matrix IN PLACE
      // This function expects to be fully gridded

      // Use elem in the function string to perform operations elementwise
      // Use val to add a variable
      
      %(GRID_IDX)s
      
      // Load blocks into shared memory
      
      __shared__ float tmp[BLOCK_SIZE][BLOCK_SIZE];

      tmp[ty][tx] = a[glob_idx]; // Switch back to row (ty), col (tx) indexing
  
      __syncthreads();

      a[glob_idx] = %(func)s;

    }
    """

def elemwise_inplace(operation, bsize = 16, val = 0, val_function = False):
    """
    Create a function which operates on each element of inp in place.
    Use 'elem' to represent each matrix element, and 'val' to represent another variable, optionally

    If val_function is true, returns a function that accepts input and a value for val. Make sure this
    is of the np.float32 dtype.

    """
    
    v = N.float32(val)

    f, block_f = Krnl(elemwise_inplace_kernel, bsize, FULL_GRID_IDX,{'func': operation}).func('elemwise_inplace', 'Pif')

    if val_function:
        return lambda inp, value: f(full_grid(inp.shape[0], inp.shape[1], bsize), block_f, inp.gpudata, inp.shape[1], value)

    return lambda inp: f(full_grid(inp.shape[0], inp.shape[1], bsize), block_f, inp.gpudata, inp.shape[1], v)


